package com.sky.wsp.license.utils;

import com.sky.wsp.license.conf.KeyToolConfig;
import org.bouncycastle.jce.X509Principal;
import org.bouncycastle.x509.X509V3CertificateGenerator;
import org.springframework.util.StringUtils;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Date;

/**
 * @author wangshaopeng@talkweb.com.cn
 * @desccription 秘钥工具类
 * @date 2024/6/24
 */
public class KeyUtil {
    public static PublicKey getPublicKey(KeyStore keyStore, KeyToolConfig keyToolConfig) {
        if (null == keyStore) {
            keyStore = loadKeyStore(keyToolConfig);
        }
        if (null == keyStore) {
            return null;
        }
        String alias = keyToolConfig.getAlias();
        try {
            if (!keyStore.containsAlias(alias)) {
                return null;
            }
            Certificate certificate = keyStore.getCertificate(alias);
            if (null == certificate) {
                return null;
            }
            return certificate.getPublicKey();
        } catch (KeyStoreException e) {
            e.printStackTrace();
            return null;
        }
    }

    public static PrivateKey getPrivateKey(KeyStore keyStore, KeyToolConfig keyToolConfig) {
        if (null == keyStore) {
            keyStore = loadKeyStore(keyToolConfig);
        }
        if (null == keyStore) {
            return null;
        }
        String alias = keyToolConfig.getAlias();
        try {
            if (!keyStore.containsAlias(alias)) {
                return null;
            }
            PrivateKey privateKey = (PrivateKey) keyStore.getKey(alias, keyToolConfig.getKeyPass().toCharArray());
            return privateKey;
        } catch (KeyStoreException | NoSuchAlgorithmException | UnrecoverableKeyException e) {
            e.printStackTrace();
            return null;
        }

    }

    public static KeyStore loadKeyStore(KeyToolConfig keyToolConfig) {
        String keystore = keyToolConfig.getKeystore();
        String storePass = keyToolConfig.getStorePass();
        if (StringUtils.isEmpty(keystore) || StringUtils.isEmpty(storePass)) {
            throw new RuntimeException("秘钥库和密码不能为空！");
        }
        String alias = keyToolConfig.getAlias();
        String password = keyToolConfig.getKeyPass();
        if (StringUtils.isEmpty(alias) || StringUtils.isEmpty(password)) {
            throw new RuntimeException("秘钥条目别名和密码不能为空！");
        }
        try {
            File ksFile = new File(keystore);
            if (ksFile.exists()) {
                KeyStore keyStore = KeyStore.getInstance("PKCS12");
                keyStore.load(new FileInputStream(ksFile), storePass.toCharArray());
                return keyStore;
            }

            // 生成RSA非对称密钥对
            KeyPair keyPair = keyPair();
            // 生成 X.509 证书（带公钥）
            X509Certificate cert = generateCertificate(keyPair);
            // RSA非对称秘钥条目
            KeyStore.PrivateKeyEntry privateKeyEntry = new KeyStore.PrivateKeyEntry(keyPair.getPrivate(), new Certificate[]{cert});
//            // AES对称秘钥
//            SecretKey secretKey = secretKey();
//            // AES对称秘钥条目
//            KeyStore.SecretKeyEntry secretKeyEntry = new KeyStore.SecretKeyEntry(secretKey);

            // 创建 PKCS12 密钥库
            KeyStore keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(null, null);

            // 将私钥和证书存入密钥库
//            keyStore.setKeyEntry(alias, keyPair.getPrivate(), password.toCharArray(), new Certificate[]{cert});
            keyStore.setEntry(alias, privateKeyEntry, new KeyStore.PasswordProtection(password.toCharArray()));
//            keyStore.setEntry(aesAlias, secretKeyEntry, new KeyStore.PasswordProtection(password.toCharArray()));

            // 保存密钥库到文件
            FileOutputStream fos = new FileOutputStream(ksFile);
            keyStore.store(fos, storePass.toCharArray());
            fos.close();
            return keyStore;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * AES对称秘钥
     */
    private static SecretKey secretKey() throws NoSuchAlgorithmException {
        KeyGenerator keygen = KeyGenerator.getInstance("AES");
        keygen.init(128);
        SecretKey aesKey = keygen.generateKey();
//        System.out.println(aesKey);
//        String encoded = Base64.encode(aesKey.getEncoded());
//        System.out.printf("aesKey %s %s %s%n", aesKey.getAlgorithm(), aesKey.getFormat(), encoded);
        return aesKey;
    }

    /**
     * RSA非对称密钥对
     */
    private static KeyPair keyPair() throws NoSuchAlgorithmException {
        // 生成密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        return keyPair;
    }

    /**
     * RSA非对称加密公钥证书
     */
    private static X509Certificate generateCertificate(KeyPair keyPair) throws Exception {
        // 生成自签名的 X.509 证书
        X509V3CertificateGenerator certGenerator = new X509V3CertificateGenerator();
        certGenerator.setSerialNumber(BigInteger.valueOf(1));
        certGenerator.setSubjectDN(new X509Principal("CN=My Certificate"));
        certGenerator.setIssuerDN(new X509Principal("CN=My Certificate"));
        certGenerator.setNotBefore(new Date());
        certGenerator.setNotAfter(new Date(System.currentTimeMillis() + 365 * 24 * 60 * 60 * 1000L));
        certGenerator.setPublicKey(keyPair.getPublic());
        certGenerator.setSignatureAlgorithm("SHA256withRSA");
        return certGenerator.generate(keyPair.getPrivate());
    }
}
